# coding: utf-8
import _thread as thread
import base64
import hashlib
import hmac
import json
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlencode, urlparse
from wsgiref.handlers import format_date_time
import websocket

answer = ""

class Ws_Param(object):
    def __init__(self, APIKey, APISecret, gpt_url):
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(gpt_url).netloc
        self.path = urlparse(gpt_url).path
        self.gpt_url = gpt_url

    def create_url(self):
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"

        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                                digestmod=hashlib.sha256).digest()
        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')

        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')

        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        return self.gpt_url + '?' + urlencode(v)

def on_error(ws, error):
    print("### error:", error)

def on_close(ws, close_status_code, close_msg):
    print("### closed ###")

def on_open(ws):
    thread.start_new_thread(run, (ws,))

def run(ws, *args):
    data = json.dumps(ws.parameter)
    ws.send(data)

def on_message(ws, message):
    global answer
    data = json.loads(message)
    code = data['header']['code']
    if code != 0:
        print(f'请求错误: {code}, {data}')
        ws.close()
        return
    
    choices = data["payload"]["choices"]
    status = choices["status"]
    content = choices["text"][0]["content"]
    answer += content
    
    if status == 2:
        ws.close()
    else:
        print(content, end='')

def main(appid, api_key, api_secret, gpt_url, domain, messages, max_tokens=5000, top_k=6, temperature=1.0):
    global answer
    answer = ""
    
    parameter = {
        "header": {
            "app_id": appid
        },
        "parameter": {
            "chat": {
                "domain": domain,
                "max_tokens": max_tokens,
                "top_k": top_k,
                "temperature": temperature,
                "show_ref_label": True
            }
        },
        "payload": {
            "message": {
                "text": messages
            }
        }
    }

    wsParam = Ws_Param(api_key, api_secret, gpt_url)
    wsUrl = wsParam.create_url()
    ws = websocket.WebSocketApp(
        wsUrl,
        on_message=on_message,
        on_error=on_error,
        on_close=on_close,
        on_open=on_open
    )
    ws.parameter = parameter
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
    return answer